#ifndef HEADER_H
#define HEADER_H

#include <time.h>
#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <string.h>
#include <math.h>
#include <sys/time.h>
#include <sys/types.h>
#include <dirent.h>
#include <cuda_runtime.h>
#include <cuda.h>
#include <vector>
#include <assert.h>
#include <algorithm>
#include <fstream>
#include <assert.h>
#include <string>
#include <ctype.h>
#include <cassert>
#include <cstdlib>
#include <sstream>

using namespace std;

// 获取一个数字的所需要的bit位数
int get_max_bit_num_of_data(unsigned int data)
{
    // 一直一维直到0
    int bit_num = 0;

    while (data != 0)
    {
        data = data >> 1;
        bit_num++;
    }

    return bit_num;
}

// 将一个bool的数组转化一个对应的unsigned int的值，与低位对齐，从高往低存
unsigned int convert_bool_vec_to_unsigned_int(vector<bool> bool_vec)
{
    int total_bit_num = bool_vec.size();
    // cout << total_bit_num << endl;
    assert(bool_vec.size() <= total_bit_num);

    unsigned int return_num = 0;

    for (int i = 0; i < bool_vec.size(); i++)
    {
        bool cur_bool = bool_vec[i];

        if (cur_bool == true)
        {
            unsigned int temp_bool_bit = 1 << (total_bit_num - 1 - i);
            return_num = return_num | temp_bool_bit;
        }
    }

    return return_num;
}

void print_coo_matrix(vector<unsigned int> row_index, vector<unsigned int> col_index, vector<float> val)
{
    assert(row_index.size() == col_index.size() && col_index.size() == val.size());

    for (int ele_id = 0; ele_id < row_index.size(); ele_id++)
    {
        cout << "(" << row_index[ele_id] << "," << col_index[ele_id] << "," << val[ele_id] << ")"
             << " , ";
    }

    cout << endl;
}

void store_coo_matrix_to_file(vector<unsigned int> row_index, vector<unsigned int> col_index, vector<float> val, string input_file)
{
    assert(row_index.size() == col_index.size() && col_index.size() == val.size());

    // 向文件中写文件
    ofstream arrWrite(input_file, ios::out | ios::trunc);

    cout << "file_name:" << input_file << endl;

    for (unsigned int i = 0; i < row_index.size(); i++)
    {
        arrWrite << row_index[i] << "," << col_index[i] << "," << val[i] << endl;
    }

    arrWrite.close();
}

// 输入一个行索引矩阵，根据行索引和线程非零元数量生成二维向量，记录每个线程要重新计算sum的非零元位置
vector<vector<bool>> get_sum_begin_bool_flag_of_each_thread(vector<unsigned int> row_index_vec, int nnz_in_thread)
{
    assert(row_index_vec.size() % nnz_in_thread == 0);

    // 一共的线程粒度的块的数量
    int TLB_num = row_index_vec.size() / nnz_in_thread;

    // 每个线程粒度的块都有一个bool标记
    vector<vector<bool>> sum_begin_bool_flag_of_each_thread(TLB_num);

    // 遍历每一个线程粒度的块，打上标记，打标记的位置有两个，一个是换行的位置，一个是warp粒度的块起始的位置。
    for (int TLB_id = 0; TLB_id < TLB_num; TLB_id++)
    {
        // 判断是不是WLB的第一个块
        bool is_first_TLB_of_WLB = false;

        if (TLB_id % 32 == 0)
        {
            is_first_TLB_of_WLB = true;
        }

        // 遍历TLB的每一个非零元
        for (int TLB_nz_id = 0; TLB_nz_id < nnz_in_thread; TLB_nz_id++)
        {
            // TLB_nz_id是0，并且TLB是WLB的第一个TLB
            if (TLB_nz_id == 0 && is_first_TLB_of_WLB == true)
            {
                sum_begin_bool_flag_of_each_thread[TLB_id].push_back(true);
            }
            else
            {
                // 非零元对应的全局行号
                unsigned int global_nz_index = TLB_id * nnz_in_thread + TLB_nz_id;
                assert(global_nz_index < row_index_vec.size());
                // 当前非零元行号对应的行号
                unsigned int row_index = row_index_vec[global_nz_index];

                assert(global_nz_index > 0);
                unsigned int previous_row_index = row_index_vec[global_nz_index - 1];

                // 如果前后的行索引一样，就输入false
                if (row_index == previous_row_index)
                {
                    sum_begin_bool_flag_of_each_thread[TLB_id].push_back(false);
                }
                else
                {
                    // 如果和之前的行索引不一样，那就可能是换行了
                    sum_begin_bool_flag_of_each_thread[TLB_id].push_back(true);
                }
            }
        }

        assert(sum_begin_bool_flag_of_each_thread[TLB_id].size() == nnz_in_thread);
    }

    return sum_begin_bool_flag_of_each_thread;
}

// 获取每个warp的首行绝对行号
vector<unsigned int> first_global_row_index_of_each_warp(vector<unsigned int> row_index_vec, int nnz_in_thread)
{
    assert(row_index_vec.size() % (32 * nnz_in_thread) == 0);

    unsigned int WLB_num = row_index_vec.size() / (32 * nnz_in_thread);

    unsigned int WLB_nz_num = 32 * nnz_in_thread;

    vector<unsigned int> WLB_first_global_row_index;

    // 遍历每一个WLB，查看WLB第一个row index，这个row index是全局行号
    for (int WLB_id = 0; WLB_id < WLB_num; WLB_id++)
    {
        // WLB第一个非零元的行号
        unsigned int WLB_first_nz_index = WLB_id * WLB_nz_num;

        // 找到对应的行号
        unsigned int WLB_first_row_index = row_index_vec[WLB_first_nz_index];

        WLB_first_global_row_index.push_back(WLB_first_row_index);
    }

    // 按照CSR的方式压缩
    WLB_first_global_row_index.push_back(row_index_vec[row_index_vec.size() - 1]);

    assert(WLB_first_global_row_index.size() == WLB_num + 1);

    

    return WLB_first_global_row_index;
}

// 获取每个线程第一个行首元素的相对行号，传入加和的起始位置，计算每个线程第一个加和起始位置的行索引，如果加和起始位置，那就按照上一个加和起始位置来
vector<unsigned int> first_relative_row_index_of_each_thread(vector<unsigned int> row_index_vec, vector<unsigned int> WLB_first_global_row_index, vector<vector<bool>> row_begin_bool_flag, int nnz_in_thread)
{
    // 查看每一个线程第一个行首元素的相对行号，主要是当前TLB之前的所有TLB的
    assert(row_index_vec.size() % (32 * nnz_in_thread) == 0);

    // TLB的数量
    unsigned int TLB_num = row_index_vec.size() / nnz_in_thread;

    vector<unsigned int> TLB_first_relative_row_index;

    for (int TLB_id = 0; TLB_id < TLB_num; TLB_id++)
    {
        // 当前TLB所在的WLB号
        unsigned int WLB_id = TLB_id / 32;

        // 当前WLB的首行行号
        unsigned int WLB_first_row_index = WLB_first_global_row_index[WLB_id];

        assert(row_begin_bool_flag[TLB_id].size() == nnz_in_thread);

        bool has_T_in_TLB = false;

        // 遍历TLB中的所有数据，找出第一个T的行号
        for (int nz_id_in_TLB = 0; nz_id_in_TLB < nnz_in_thread; nz_id_in_TLB++)
        {
            if (row_begin_bool_flag[TLB_id][nz_id_in_TLB] == true)
            {
                has_T_in_TLB = true;
                unsigned int TLB_global_first_row_index = row_index_vec[TLB_id * nnz_in_thread + nz_id_in_TLB];

                assert(TLB_global_first_row_index >= WLB_first_row_index);
                unsigned int TLB_relative_first_row_index = TLB_global_first_row_index - WLB_first_row_index;
                // 写到对应位置
                TLB_first_relative_row_index.push_back(TLB_relative_first_row_index);
                break;
            }
        }

        // 如果全是F，就继承之前的行号
        if (has_T_in_TLB == false)
        {
            assert(TLB_first_relative_row_index.size() > 0);
            TLB_first_relative_row_index.push_back(TLB_first_relative_row_index[TLB_first_relative_row_index.size() - 1]);
        }
    }

    return TLB_first_relative_row_index;
}

// 获取归约数组，当行比较长的之后，线程边界把一行分为一个头，多个身子，和一个尾巴。头是一定存在ad，身子不一定存在，尾巴在逻辑上存在（默认为0）
// 首先要先将身子和尾巴执行一次归约，然后再执行脑袋和其他部分的规约
// 而有多少个身子情况是复杂的，身子的数量决定了身子和尾巴的归约过程。因为行脑袋和行其他部分的归约结果在相邻的两个线程，所以会先把行身子和行尾巴的中间
// 结果先向前挪一个线程，然后再执行规约。在同一个线程内执行脑袋和其他部分的拼装。这也决定了行身子和行尾巴的归约最终是在行脑袋所在的线程中进行的。
// 在行身子和行尾巴的归约中，需要记录行身子数量，从而记录归约的长度。而行身子的数量要存在行头部所在的线程中。
// 所以最终要做的是，将bool flag全是F的相邻线程粒度的块的数量记录下来，存在最前面的带T的线程粒度的块对应的位置上
// 
vector<unsigned int> get_reduce_offset_vec(vector<vector<bool>> begin_sum_bool_flag)
{
    assert(begin_sum_bool_flag.size() > 0);

    // 第一个块肯定有T
    {
        bool has_T_in_first_thread = false;

        for (int i = 0; i < begin_sum_bool_flag[0].size(); i++)
        {
            has_T_in_first_thread = has_T_in_first_thread || begin_sum_bool_flag[0][i];
        }

        assert(has_T_in_first_thread == true);
    }

    // 上一个有T的线程id
    int last_has_T_TLB_id = 0;
    // 累计全F的行的数量
    int TLB_num_with_all_F = 0;

    vector<unsigned int> reduce_offset_vec(begin_sum_bool_flag.size());

    // 遍历所有的线程
    for (int TLB_id = 1; TLB_id < begin_sum_bool_flag.size(); TLB_id++)
    {
        bool has_T = false;
        // 首先判断是不是全是F
        for (int nz_id_in_TLB = 0; nz_id_in_TLB < begin_sum_bool_flag[TLB_id].size(); nz_id_in_TLB++)
        {
            has_T = has_T || begin_sum_bool_flag[TLB_id][nz_id_in_TLB];
        }

        // 如果没有T就累加了一下
        if (has_T == false)
        {
            TLB_num_with_all_F++;
        }
        else
        {
            // 到了有T的，记录归约的位置
            assert(last_has_T_TLB_id < begin_sum_bool_flag.size());
            reduce_offset_vec[last_has_T_TLB_id] = TLB_num_with_all_F;
            TLB_num_with_all_F = 0;
            last_has_T_TLB_id = TLB_id;
        }
    }

    // 最后一个TLB的
    assert(last_has_T_TLB_id < begin_sum_bool_flag.size());
    reduce_offset_vec[last_has_T_TLB_id] = TLB_num_with_all_F;

    // 返回对应的归约偏移量
    return reduce_offset_vec;
}

// 将bool bit、归约偏移、相对行号合并到一个数据中，高位是bool flag产生的bit flag。低位是紧挨着的归约偏移和相对行号
// 高位|bit_flag|可能存在的空白|归约偏移|相对行号|低位
// 三个数据所占的bit数量
vector<unsigned int> combine_meta_to_a_unsigned_int_vec(vector<vector<bool>> bool_flag, vector<unsigned int> reduce_offset, vector<unsigned int> relative_row_index, int *bit_num_of_bit_flag, int *bit_num_of_reduce_offset, int *bit_num_of_relative_row_index)
{
    assert(bit_num_of_bit_flag != NULL && bit_num_of_reduce_offset != NULL && bit_num_of_relative_row_index != NULL);
    assert(bool_flag.size() == reduce_offset.size() && reduce_offset.size() == relative_row_index.size());

    // 找到reduce_offset的最大值
    unsigned int max_of_reduce_offset = *max_element(reduce_offset.begin(), reduce_offset.end());
    unsigned int max_of_relative_row_index = *max_element(relative_row_index.begin(), relative_row_index.end());
    int total_bit_size = sizeof(unsigned int) * 8;

    *bit_num_of_reduce_offset = get_max_bit_num_of_data(max_of_reduce_offset);
    *bit_num_of_relative_row_index = get_max_bit_num_of_data(max_of_relative_row_index);
    *bit_num_of_bit_flag = bool_flag[0].size();

    // cout << "*bit_num_of_reduce_offset:" << *bit_num_of_reduce_offset << endl;
    // cout << "*bit_num_of_relative_row_index:" << *bit_num_of_relative_row_index << endl;
    // cout << "*bit_num_of_bit_flag:" << *bit_num_of_bit_flag << endl;

    // 空白部分的大小
    int bit_num_empty_bit = total_bit_size - *bit_num_of_reduce_offset - *bit_num_of_relative_row_index - *bit_num_of_bit_flag;

    vector<unsigned int> combine_meta_vec(reduce_offset.size());

    assert(*bit_num_of_reduce_offset + *bit_num_of_relative_row_index + *bit_num_of_bit_flag <= total_bit_size);

    // 从bit_flag开始
    for (unsigned int combine_meta_vec_id = 0; combine_meta_vec_id < combine_meta_vec.size(); combine_meta_vec_id++)
    {
        unsigned int cur_combine_meta = 0;
        
        // for (int i = 0; i < bool_flag[combine_meta_vec_id].size(); i++)
        // {
        //     cout << bool_flag[combine_meta_vec_id][i] << ",";
        // }

        // cout << endl;

        // 首先将bool flag填入
        unsigned int bit_flag = convert_bool_vec_to_unsigned_int(bool_flag[combine_meta_vec_id]);

        // 先把bit_flag添加进去
        cur_combine_meta = cur_combine_meta | bit_flag;

        // cout << "cur_combine_meta:" << cur_combine_meta << " bit_flag:" << bit_flag << endl;

        // exit(-1);

        // 挪一个空白和归约偏移的距离
        cur_combine_meta = cur_combine_meta << (bit_num_empty_bit + *bit_num_of_reduce_offset);
        
        // 将归约偏移放进来
        cur_combine_meta = cur_combine_meta | reduce_offset[combine_meta_vec_id];

        // 空出来一个相对行号的大小
        cur_combine_meta = cur_combine_meta << (*bit_num_of_relative_row_index);

        cur_combine_meta = cur_combine_meta | relative_row_index[combine_meta_vec_id];

        // 将数据插入到
        combine_meta_vec[combine_meta_vec_id] = cur_combine_meta;
    }
    
    return combine_meta_vec;
}

// 打印一个数组
void print_vec(vector<unsigned int> vec)
{
    // 打印对应的数组
    for (int item_id = 0; item_id < vec.size(); item_id++)
    {
        cout << "(" << item_id << "," << vec[item_id] << ") , ";
    }

    cout << endl;
}

// 将数组打印到文件中
template<typename T>
void store_vec_to_file(vector<T> vec, string input_file)
{
    // 向文件中写文件
    ofstream arrWrite(input_file, ios::out | ios::trunc);

    cout << "file_name:" << input_file << endl;

    for (unsigned int i = 0; i < vec.size(); i++)
    {
        arrWrite << vec[i] << endl;
    }

    arrWrite.close();
}


// 将合并的元数据按照一定bit
void store_combine_meta_vec(vector<unsigned int> combine_meta_vec, string input_file)
{
    // 向文件中写文件
    ofstream arrWrite(input_file, ios::out | ios::trunc);

    cout << "file_name:" << input_file << endl;

    int combine_meta_total_bit_num = sizeof(unsigned int) * 8;

    for (unsigned int i = 0; i < combine_meta_vec.size(); i++)
    {
        unsigned int combine_meta = combine_meta_vec[i];

        for (int bit_id = 0; bit_id < combine_meta_total_bit_num; bit_id++)
        {
            // 遍历每一个bit，从高位开始遍历，移位，取第一位，然后打印
            unsigned int bit = (combine_meta >> (combine_meta_total_bit_num - 1 - bit_id)) & 0x1;
            
            arrWrite << bit;

            if ((bit_id + 1) % 4 == 0)
            {
                arrWrite << " ";
            }
        }

        arrWrite << endl;
    }

    arrWrite.close();
}

// 打印加和起始位置的二维数组
void store_bool_flag_of_sum_begin_to_file(vector<vector<bool>> bool_vec, string input_file)
{
    // 向文件中写文件
    ofstream arrWrite(input_file, ios::out | ios::trunc);

    cout << "file_name:" << input_file << endl;

    // 打印
    for (unsigned int i = 0; i < bool_vec.size(); i++)
    {
        for (unsigned int j = 0; j < bool_vec[i].size(); j++)
        {
            arrWrite << bool_vec[i][j] << " , ";
        }
        arrWrite << endl;
    }

    arrWrite.close();
}

// 执行warp级别的交错操作，将warp内的非零元根据WLB内的TLB号交换其未知
vector<unsigned int> stagger_col_index_in_warp(vector<unsigned int> col_index_vec)
{
    // 一开始先不去实现
    return col_index_vec;
}

vector<float> stagger_val_in_warp(vector<float> val_vec)
{
    // 一开始先不去实现
    return val_vec;
}

void split(const std::string &s, std::vector<std::string> &sv, const char delim = ' ')
{
    sv.clear();
    std::istringstream iss(s);
    std::string temp;

    while (std::getline(iss, temp, delim))
    {
        sv.emplace_back(std::move(temp));
    }

    return;
}

// 从外部读入一个矩阵文件，头部直接是矩阵规模，矩阵是先按照行，然后按照列排序好。并且没有矩阵，所以需要一个预处理
// 自动对空行执行一个非零元的padding
void get_coo_vector_of_matrix_from_file(string file_name, unsigned int* col_num, unsigned int* row_num, unsigned int* nnz, vector<unsigned int> &row_index_arr, vector<unsigned int> &col_index_arr, vector<float> &val_arr)
{
    assert(row_index_arr.size() == 0 && col_index_arr.size() == 0 && val_arr.size() == 0);
    // 从对应文件读内容
    // 读文件
    char buf[1024];

    ifstream infile;
    infile.open(file_name);

    bool dataset_first_line = true;

    // 被遍历到的最后一行行号
    unsigned long last_row_index = 0;

    if (infile.is_open())
    {
        // 读弗洛里达矩阵格式，第一行是矩阵规模，先是行数量，然后是列数量
        while (infile.good() && !infile.eof())
        {
            string line_str;
            vector<string> sv;
            memset(buf, 0, 1024);
            infile.getline(buf, 1024);
            line_str = buf;

            // 碰到奇怪的输入就跳过
            if (isspace(line_str[0]) || line_str.empty())
            {
                continue;
            }

            split(line_str, sv);

            // 矩阵的规模，先是最大行号，然后是最大列号
            if (dataset_first_line == true)
            {
                // 矩阵的规模，因为索引从0开始，所以索引的取值范围应该-1
                *row_num = atol(sv[0].c_str());
                *col_num = atol(sv[1].c_str());
                *nnz = atol(sv[2].c_str());
                dataset_first_line = false;
                continue;
            }

            // 佛罗里达矩阵先行索引，然后是列索引
            unsigned long row_index = atol(sv[0].c_str()) - 1;
            unsigned long col_index = atol(sv[1].c_str()) - 1;

            assert(row_index >= last_row_index);
            // 这里判断是不是出现了空行
            if (row_index > last_row_index + 1)
            {
                // 出现了空行，将空行padding一下。
                for (unsigned int row_id = last_row_index + 1; row_id < row_index; row_id++)
                {
                    row_index_arr.push_back(row_id);
                    col_index_arr.push_back(0);
                    val_arr.push_back(0);

                    *nnz = *nnz + 1;
                }
            }

            last_row_index = row_index;

            // 分别是行、列、值
            row_index_arr.push_back(row_index);
            col_index_arr.push_back(col_index);

            // val_arr.push_back(stof(sv[2].c_str()));
            val_arr.push_back(1);
        }
    }
    
    // cout << row_index_arr.size() << " " << *nnz << endl;
    assert(row_index_arr.size() == *nnz);
    assert(row_index_arr[row_index_arr.size() - 1] == *row_num - 1);

    infile.close();
}

// 将nnz padding到对应位置
void padding_coo_matrix_to_target_nnz(vector<unsigned int> &row_index_arr, vector<unsigned int> &col_index_arr, vector<float> &val_arr, unsigned int target_nnz)
{
    assert(target_nnz > row_index_arr.size());

    assert(row_index_arr.size() == col_index_arr.size() && val_arr.size() == row_index_arr.size());
    
    // 要padding的数量
    unsigned int adden_nz = target_nnz - row_index_arr.size();

    // 针对最后一行执行padding，比较方便
    for (unsigned int i = 0; i < adden_nz; i++)
    {
        row_index_arr.push_back(row_index_arr[row_index_arr.size() - 1]);
        col_index_arr.push_back(col_index_arr[col_index_arr.size() - 1]);
        val_arr.push_back(0);
    }

    assert(row_index_arr.size() == target_nnz);
}


#endif
